// Copyright 2021 The TensorStore Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <stdint.h>

#include <string_view>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/cord.h"
#include "riegeli/bytes/cord_reader.h"
#include "riegeli/bytes/cord_writer.h"
#include "riegeli/bytes/string_reader.h"
#include "tensorstore/internal/image/avif_reader.h"
#include "tensorstore/internal/image/avif_writer.h"
#include "tensorstore/internal/image/image_info.h"
#include "tensorstore/util/span.h"
#include "tensorstore/util/status_testutil.h"

namespace {

using ::tensorstore::IsOk;
using ::tensorstore::StatusIs;
using ::tensorstore::internal_image::AvifReader;
using ::tensorstore::internal_image::AvifReaderOptions;
using ::tensorstore::internal_image::AvifWriter;
using ::tensorstore::internal_image::AvifWriterOptions;
using ::tensorstore::internal_image::ImageInfo;

TEST(AvifTest, Decode) {
  // https://shoonia.github.io/1x1/#405060ff
  static constexpr unsigned char data[] = {
      0x00, 0x00, 0x00, 0x20, 0x66, 0x74, 0x79, 0x70, 0x61, 0x76, 0x69, 0x66,
      0x00, 0x00, 0x00, 0x00, 0x61, 0x76, 0x69, 0x66, 0x6d, 0x69, 0x66, 0x31,
      0x6d, 0x69, 0x61, 0x66, 0x4d, 0x41, 0x31, 0x41, 0x00, 0x00, 0x00, 0xf2,
      0x6d, 0x65, 0x74, 0x61, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x28,
      0x68, 0x64, 0x6c, 0x72, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x70, 0x69, 0x63, 0x74, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x6c, 0x69, 0x62, 0x61, 0x76, 0x69, 0x66, 0x00,
      0x00, 0x00, 0x00, 0x0e, 0x70, 0x69, 0x74, 0x6d, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x01, 0x00, 0x00, 0x00, 0x1e, 0x69, 0x6c, 0x6f, 0x63, 0x00, 0x00,
      0x00, 0x00, 0x44, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
      0x00, 0x00, 0x01, 0x1a, 0x00, 0x00, 0x00, 0x1f, 0x00, 0x00, 0x00, 0x28,
      0x69, 0x69, 0x6e, 0x66, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,
      0x00, 0x1a, 0x69, 0x6e, 0x66, 0x65, 0x02, 0x00, 0x00, 0x00, 0x00, 0x01,
      0x00, 0x00, 0x61, 0x76, 0x30, 0x31, 0x43, 0x6f, 0x6c, 0x6f, 0x72, 0x00,
      0x00, 0x00, 0x00, 0x6a, 0x69, 0x70, 0x72, 0x70, 0x00, 0x00, 0x00, 0x4b,
      0x69, 0x70, 0x63, 0x6f, 0x00, 0x00, 0x00, 0x14, 0x69, 0x73, 0x70, 0x65,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
      0x00, 0x00, 0x00, 0x10, 0x70, 0x69, 0x78, 0x69, 0x00, 0x00, 0x00, 0x00,
      0x03, 0x08, 0x08, 0x08, 0x00, 0x00, 0x00, 0x0c, 0x61, 0x76, 0x31, 0x43,
      0x81, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x13, 0x63, 0x6f, 0x6c, 0x72,
      0x6e, 0x63, 0x6c, 0x78, 0x00, 0x01, 0x00, 0x0d, 0x00, 0x00, 0x80, 0x00,
      0x00, 0x00, 0x17, 0x69, 0x70, 0x6d, 0x61, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x01, 0x00, 0x01, 0x04, 0x01, 0x02, 0x83, 0x04, 0x00, 0x00,
      0x00, 0x27, 0x6d, 0x64, 0x61, 0x74, 0x12, 0x00, 0x0a, 0x07, 0x38, 0x00,
      0x06, 0x10, 0x10, 0xd0, 0x02, 0x32, 0x12, 0x10, 0x00, 0x00, 0x00, 0x0f,
      0xfa, 0x40, 0x1a, 0xa8, 0x31, 0xea, 0x80, 0x78, 0x73, 0xf9, 0x11, 0x7b,
      0x50,
  };

  AvifReader decoder;
  riegeli::StringReader string_reader(reinterpret_cast<const char*>(data),
                                      sizeof(data));
  ASSERT_THAT(decoder.Initialize(&string_reader), IsOk());

  const auto info = decoder.GetImageInfo();
  EXPECT_EQ((ImageInfo{1, 1, 3}), info);

  uint8_t pixel[3] = {};
  ASSERT_THAT(decoder.Decode(pixel), IsOk());

  EXPECT_EQ(0x40, pixel[0]);
  EXPECT_EQ(0x50, pixel[1]);
  EXPECT_EQ(0x60, pixel[2]);
}

TEST(AvifTest, EncodeDecode) {
  // https://aomediacodec.github.io/av1-avif/#baseline-profile
  // The baseline profile for avif uses ftype=avif and
  // compatible_brands avif,mif1,miaf.
  // clang-format off
  static constexpr const char prefix[] = {
      0x0, 0x0, 0x0, 0x1c, 'f', 't', 'y', 'p', 'a', 'v', 'i', 'f',
      0x0, 0x0, 0x0, 0x0,  'a', 'v', 'i', 'f', 'm', 'i', 'f', '1',
      'm', 'i', 'a', 'f',
  };
  // clang-format on

  uint8_t pixels[1] = {42};

  absl::Cord encoded;
  {
    AvifWriter encoder;
    riegeli::CordWriter cord_writer(&encoded);
    ASSERT_THAT(encoder.Initialize(&cord_writer), IsOk());

    ASSERT_THAT(encoder.Encode(ImageInfo{1, 1, 1}, pixels), IsOk());
    ASSERT_THAT(encoder.Done(), IsOk());
  }

  EXPECT_THAT(encoded.Flatten(),
              ::testing::StartsWith(std::string_view(prefix, sizeof(prefix))));

  {
    AvifReader decoder;
    riegeli::CordReader cord_reader(&encoded);
    ASSERT_THAT(decoder.Initialize(&cord_reader), IsOk());

    const auto& info = decoder.GetImageInfo();
    EXPECT_EQ((ImageInfo{1, 1, 1}), info);

    uint8_t new_pixels[1] = {};
    ASSERT_THAT(decoder.Decode(new_pixels), IsOk());
    EXPECT_NEAR(new_pixels[0], pixels[0], 4);
  }
}

TEST(AvifTest, CorruptData) {
  static constexpr unsigned char data[] = {
      0x00, 0x00, 0x00, 0x1c, 0x66, 0x74, 0x79, 0x70, 0x6d, 0x69, 0x66, 0x31,
      0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};

  riegeli::StringReader string_reader(reinterpret_cast<const char*>(data),
                                      sizeof(data));

  AvifReader decoder;
  EXPECT_THAT(decoder.Initialize(&string_reader),
              StatusIs(absl::StatusCode::kInvalidArgument));
}

}  // namespace
